#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void binary_cross_entropy_custom_do(uint32_t blockDim, void *stream, uint8_t *p, uint8_t *y, uint8_t *z);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void binary_cross_entropy_custom(GM_ADDR p, GM_ADDR y, GM_ADDR z);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t ByteSize = 8 * 2048 * sizeof(uint16_t);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *p = (uint8_t *)AscendC::GmAlloc(ByteSize);
    uint8_t *y = (uint8_t *)AscendC::GmAlloc(ByteSize);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(ByteSize);

    ReadFile("./input/input_p.bin", ByteSize, p, ByteSize);
    ReadFile("./input/input_y.bin", ByteSize, y, ByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(binary_cross_entropy_custom, blockDim, p, y, z); // use this macro for cpu debug

    WriteFile("./output/output_z.bin", z, ByteSize);

    AscendC::GmFree((void *)p);
    AscendC::GmFree((void *)y);
    AscendC::GmFree((void *)z);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *pHost, *yHost, *zHost;
    uint8_t *pDevice, *yDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&pHost), ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), ByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), ByteSize));
    CHECK_ACL(aclrtMalloc((void **)&pDevice, ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, ByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_p.bin", ByteSize, pHost, ByteSize);
    ReadFile("./input/input_y.bin", ByteSize, yHost, ByteSize);

    CHECK_ACL(aclrtMemcpy(pDevice, ByteSize, pHost, ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(yDevice, ByteSize, yHost, ByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    
    binary_cross_entropy_custom_do(blockDim, stream, pDevice, yDevice, zDevice);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, ByteSize, zDevice, ByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, ByteSize);

    CHECK_ACL(aclrtFree(pDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(pHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(zHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
